import torch
import torchvision

model = torchvision.models.vgg16()
# print(model)

# print(torch.cuda.is_available())
# print(torch.cuda.device_count())
# print(torch.cuda.get_device_name(0))
# print(torch.cuda.current_device())

# device = torch.device("cuda:0")
# print(device)

# 根据设备情况自动选择cpu或者gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# 将模型放在gpu上
# model.to(device)
# 或者
model.cuda()
data = torch.ones((2,3))
print(data)
# 将数据放在gpu
data = data.to(device)
data = data.cuda()
data = data.cpu()